import numpy as np


# this class add noise to data
class Noise:
    def __init__(self):
        pass

    # this function add laplace noise to some data, real data is object_array. Sensitivity of
    # laplace mechanism is sensitivity1 and privacy budget is epsilon1.
    def add_laplace(self, object_array, epsilon1, sensitivity1, if_regularize=True):
        lambda1 = sensitivity1 / epsilon1
        shape1 = object_array.shape
        noisy_object = object_array + np.random.laplace(scale=lambda1, size=shape1)
        if if_regularize:
            noisy_object = self.positive_regulation(noisy_object)
        return noisy_object

    def add_laplace_for_tree(self, object_array, epsilon1, sensitivity1, if_regularize=True):
        lambda1 = sensitivity1 / epsilon1
        shape1 = object_array.shape
        noise = np.random.laplace(scale=lambda1, size=shape1)
        noisy_object = object_array + noise
        # temp = noisy_object
        if if_regularize:
            noisy_object = self.positive_regulation_for_tree(noisy_object)
        # if np.allclose(noisy_object, 0):
        #     print(f"加的噪声为: {noise}")
        #     print(f"未规则化前:{temp}")
        return noisy_object

    # this function change meaningless negative value to zero. These negative value comes from
    # noise adding and usually meaningless such as density. This function sum up all negative value
    # and add this negative sum to the smallest positive value. Iter this until the negative sum runs out.
    # 此函数将无意义的负值更改为零。这些负值来自
    # 噪声增加，通常没有意义，如密度。此函数对所有负值进行求和
    # 并将该负和与最小正值相加。一直这样，直到负和用完。
    def positive_regulation(self, noisy_array1):
        sort_indices = np.argsort(noisy_array1)
        sorted_array1 = noisy_array1[sort_indices]
        negative_indices = np.argwhere(sorted_array1 < 0).reshape(-1)
        positive_indices = np.argwhere(sorted_array1 > 0).reshape(-1)

        if negative_indices.size > 0:
            if positive_indices.size > 0:
                negative_sum = np.sum(sorted_array1[negative_indices])
                positive_sum = np.sum(sorted_array1[positive_indices])
                if positive_sum > - negative_sum:
                    index_positive_larger_than_negative = 1
                    temp_positive_sum = np.sum(sorted_array1[positive_indices[0:index_positive_larger_than_negative]])#切片包头不包尾
                    while - negative_sum > temp_positive_sum:
                        index_positive_larger_than_negative += 1
                        temp_positive_sum = np.sum(sorted_array1[positive_indices[0:index_positive_larger_than_negative]])
                    sorted_array1[positive_indices[index_positive_larger_than_negative - 1]] = np.sum(
                        sorted_array1[positive_indices[0:index_positive_larger_than_negative]]) + negative_sum
                    if index_positive_larger_than_negative > 1:
                        sorted_array1[positive_indices[0:index_positive_larger_than_negative - 1]] = 0
                    sorted_array1[negative_indices] = 0
                    indices_inverse = np.argsort(sort_indices)
                    array_after_regulation = sorted_array1[indices_inverse]#将数组还原成排序前的位置
                else:
                    array_after_regulation = np.zeros(noisy_array1.size)
            else:
                array_after_regulation = np.zeros(noisy_array1.size)
        else:
            array_after_regulation = noisy_array1
        return array_after_regulation

    def positive_regulation_for_tree(self, noisy_array1):
        sort_indices = np.argsort(noisy_array1)
        sorted_array1 = noisy_array1[sort_indices]
        negative_indices = np.argwhere(sorted_array1 < 0).reshape(-1)
        positive_indices = np.argwhere(sorted_array1 > 0).reshape(-1)

        if negative_indices.size > 0:
            if positive_indices.size > 0:
                negative_sum = np.sum(sorted_array1[negative_indices])
                positive_sum = np.sum(sorted_array1[positive_indices])
                if positive_sum > - negative_sum:
                    index_positive_larger_than_negative = 1
                    temp_positive_sum = np.sum(sorted_array1[positive_indices[0:index_positive_larger_than_negative]])#切片包头不包尾
                    while - negative_sum > temp_positive_sum:
                        index_positive_larger_than_negative += 1
                        temp_positive_sum = np.sum(sorted_array1[positive_indices[0:index_positive_larger_than_negative]])
                    sorted_array1[positive_indices[index_positive_larger_than_negative - 1]] = np.sum(
                        sorted_array1[positive_indices[0:index_positive_larger_than_negative]]) + negative_sum
                    if index_positive_larger_than_negative > 1:
                        sorted_array1[positive_indices[0:index_positive_larger_than_negative - 1]] = 0
                    sorted_array1[negative_indices] = 0
                    indices_inverse = np.argsort(sort_indices)
                    array_after_regulation = sorted_array1[indices_inverse]#将数组还原成排序前的位置
                else:
                    array_after_regulation = np.zeros(noisy_array1.size)
            else:
                array_after_regulation = np.zeros(noisy_array1.size)
        else:
            array_after_regulation = noisy_array1
        return array_after_regulation

    # this function regulate matrix after adding noise by row
    def positive_regulation_for_markov_matrix(self, noisy_markov, regulation_method='queue_minus'):
        if regulation_method == 'queue_minus':
            for rows in range(noisy_markov.shape[0]):
                noisy_markov[rows] = self.positive_regulation(noisy_markov[rows])
        elif regulation_method == 'truncation':
            noisy_markov[noisy_markov < 0] = 0
        return noisy_markov
